import numpy as np
import math


class GridFiller:
    def __init__(self, grid_size_or_preset_grid, cell_size):
        # Check if the first argument is a grid size tuple or a preset grid
        if isinstance(grid_size_or_preset_grid, tuple):
            # Initialize with grid size
            self.grid_size = (math.ceil(grid_size_or_preset_grid[0] / cell_size),
                              math.ceil(grid_size_or_preset_grid[1] / cell_size))
            self.grid = np.zeros(self.grid_size)
        else:
            # Initialize with preset grid
            self.grid = grid_size_or_preset_grid.copy()
            self.grid_size = self.grid.shape

        self.cell_size = cell_size

    @staticmethod
    def point_in_polygon(x, y, vertices):
        n = len(vertices)
        inside = False

        p1x, p1y = vertices[0]
        for i in range(1, n + 1):
            p2x, p2y = vertices[i % n]
            if min(p1y, p2y) < y <= max(p1y, p2y) and x <= max(p1x, p2x):
                if p1y != p2y and (p1x == p2x or x <= p1x + (y - p1y) * (p2x - p1x) / (p2y - p1y)):
                    inside = not inside
            p1x, p1y = p2x, p2y

        return inside

    def fill_polygon(self, vertices):
        # Convert input vertices to grid coordinates
        grid_vertices = [(int(x / self.cell_size), int(y / self.cell_size)) for x, y in vertices]

        if len(grid_vertices) == 1:
            self.fill_point(grid_vertices)
        elif len(grid_vertices) == 2:
            self.fill_line(grid_vertices)
        else:
            self.fill_shape(grid_vertices)

    def fill_line(self, vertices):
        (x0, y0), (x1, y1) = vertices
        steep = abs(y1 - y0) > abs(x1 - x0)

        if steep:
            x0, y0 = y0, x0
            x1, y1 = y1, x1

        if x0 > x1:
            x0, x1 = x1, x0
            y0, y1 = y1, y0

        dx = x1 - x0
        dy = abs(y1 - y0)
        error = int(dx / 2)
        ystep = 1 if y0 < y1 else -1
        y = y0

        for x in range(x0, x1 + 1):
            if 0 <= x < self.grid_size[1] and 0 <= y < self.grid_size[0] and steep:
                self.grid[y, x] = 1
            elif 0 <= x < self.grid_size[0] and 0 <= y < self.grid_size[1] and not steep:
                self.grid[x, y] = 1
            error -= dy
            if error < 0:
                y += ystep
                error += dx

    def fill_point(self, vertices):
        x, y = vertices[0]
        x, y = int(x), int(y)
        if 0 <= x < self.grid_size[0] and 0 <= y < self.grid_size[1]:
            self.grid[x, y] = 1

    def fill_shape(self, vertices):
        for j in range(self.grid_size[0]):
            for i in range(self.grid_size[1]):
                x = j + 0.5
                y = i + 0.5
                if self.point_in_polygon(x, y, vertices):
                    self.grid[j, i] = 1

        n = len(vertices)
        for i in range(n):
            self.fill_line([vertices[i], vertices[(i + 1) % n]])

    def fill_circle(self, circle_center, circle_radius):
        # Convert input circle_center to grid coordinates
        grid_circle_center = (
            circle_center[0] / self.cell_size, circle_center[1] / self.cell_size)
        grid_circle_radius = math.ceil(circle_radius / self.cell_size)

        for j in range(self.grid_size[0]):
            for i in range(self.grid_size[1]):
                x = j + 0.5
                y = i + 0.5
                if (x - grid_circle_center[0]) ** 2 + (y - grid_circle_center[1]) ** 2 <= grid_circle_radius ** 2:
                    self.grid[j, i] = 1

    def fill_ellipse(self, F1, F2, short_axis):
        # Convert input foci to grid coordinates
        grid_F1 = (F1[0] / self.cell_size, F1[1] / self.cell_size)
        grid_F2 = (F2[0] / self.cell_size, F2[1] / self.cell_size)

        # Compute half the distance between the foci
        c = np.sqrt((grid_F1[0] - grid_F2[0]) ** 2 + (grid_F1[1] - grid_F2[1]) ** 2) / 2

        # Compute long_axis using the given short_axis
        long_axis = np.sqrt(c ** 2 + (short_axis / self.cell_size / 2) ** 2)

        # Compute the center of the ellipse
        center = ((grid_F1[0] + grid_F2[0]) / 2, (grid_F1[1] + grid_F2[1]) / 2)

        # Compute the rotation angle of the ellipse
        theta = np.arctan2(grid_F2[1] - grid_F1[1], grid_F2[0] - grid_F1[0])

        for j in range(self.grid_size[0]):
            for i in range(self.grid_size[1]):
                x = j + 0.5
                y = i + 0.5

                # Transform the grid point to the ellipse's local coordinate system
                x_local = (x - center[0]) * np.cos(theta) + (y - center[1]) * np.sin(theta)
                y_local = -(x - center[0]) * np.sin(theta) + (y - center[1]) * np.cos(theta)

                # Check if the transformed point is inside the ellipse, with a small offset for grid points that touch the ellipse border
                if (x_local / (long_axis + 0.5)) ** 2 + (y_local / (short_axis / self.cell_size / 2 + 0.5)) ** 2 <= 1:
                    self.grid[j, i] = 1

    def dilate_obstacles(self, dilation_size):
        grid_size = self.grid.shape
        dilated_grid = np.zeros(grid_size)

        # Convert dilation_size to grid coordinates and round up
        grid_dilation_size = math.ceil(dilation_size / self.cell_size)

        for j in range(grid_size[0]):
            for i in range(grid_size[1]):
                if self.grid[j, i] == 1:
                    min_row = max(0, j - grid_dilation_size)
                    max_row = min(grid_size[0], j + grid_dilation_size + 1)
                    min_col = max(0, i - grid_dilation_size)
                    max_col = min(grid_size[1], i + grid_dilation_size + 1)

                    dilated_grid[min_row:max_row, min_col:max_col] = 1
            print(f"正在膨胀处理第{j}行")

        return dilated_grid
